#ifndef ARMOR_DETECTOR__NUMBER_CLASSIFIER_HPP_
#define ARMOR_DETECTOR__NUMBER_CLASSIFIER_HPP_

#include <openvino/openvino.hpp>
#include <opencv2/opencv.hpp>
#include <rclcpp/logger.hpp>

#include <memory>
#include <string>
#include <vector>

#include "armor.hpp"

namespace rm_aim
{

class NumberClassifier
{
public:
  /**
   * @brief 构造函数
   * @param model_path 模型文件路径
   * @param label_path 标签文件路径
   * @param threshold 分类置信度阈值
   * @param logger ROS2日志对象
   * @param ignore_classes 需要忽略的类别
   */
  NumberClassifier(
    const std::string & model_path,
    const std::string & label_path,
    const double threshold,
    const rclcpp::Logger & logger,
    const std::vector<std::string> & ignore_classes = {});

  /**
   * @brief 在指定设备上编译模型
   * @param device 设备名称 (CPU, GPU, MYRIAD等)
   */
  void compileModel(const std::string & device = "CPU");

  /**
   * @brief 从装甲板中提取数字区域
   * @param src 原始图像
   * @param armors 装甲板列表
   */
  void extractNumbers(const cv::Mat & src, std::vector<Armor> & armors);

  /**
   * @brief 对提取的数字进行分类并过滤无效装甲板
   * @param armors 装甲板列表
   */
  void classify(std::vector<Armor> & armors);

private:
  /**
   * @brief 加载标签文件
   * @param label_path 标签文件路径
   */
  void loadLabels(const std::string & label_path);

  /**
   * @brief 准备输入张量
   * @param image 输入图像
   * @return 准备好的OpenVINO张量
   */
  ov::Tensor prepareInputTensor(const cv::Mat & image);

  /**
   * @brief 处理模型输出并更新装甲板信息
   * @param armor 装甲板对象
   */
  void processOutput(Armor & armor);

  /**
   * @brief 计算softmax激活函数
   * @param input 输入数据
   * @param size 数据大小
   * @param output 输出结果
   */
  void computeSoftmax(const float * input, size_t size, float * output);

  /**
   * @brief 过滤无效装甲板
   * @param armors 装甲板列表
   */
  void filterArmors(std::vector<Armor> & armors);

  double threshold_;                          // 置信度阈值
  rclcpp::Logger logger_;                     // ROS2日志对象
  std::vector<std::string> class_names_;      // 类别名称列表
  std::vector<std::string> ignore_classes_;   // 需忽略的类别

  std::unique_ptr<ov::Core> core_;            // OpenVINO核心对象
  std::shared_ptr<ov::Model> model_;          // 模型对象
  ov::CompiledModel compiled_model_;          // 编译后的模型
  ov::InferRequest infer_request_;            // 推理请求对象
};

}  // namespace rm_aim

#endif  // ARMOR_DETECTOR__NUMBER_CLASSIFIER_HPP_